import argparse


def get_args():
    parser = argparse.ArgumentParser()
    # Required parameters
    parser.add_argument("--datasets",
                        nargs='+',
                        default=['cub', 'dogs', 'airs'])
    parser.add_argument("--imgs_path",
                        nargs='+',
                        default=[r'E:\desktop\data\dogs\Images\n02085620-Chihuahua\n02085620_199.jpg'])
    parser.add_argument("--models",
                        nargs='+',
                        default=['resnet50', 'vit-b-16','vit-l-16'])

    ################ draw loss ####################
    parser.add_argument("--loss_path",
                        type=str,
                        default=r'E:\desktop\毕业论文\代码')

    ################ draw cam #####################
    parser.add_argument("--resnet_path",
                        nargs='*',
                        default=[r'E:\desktop\毕业论文\代码\resnet50\dogs\resnet50_weights.pth'])
    parser.add_argument("--vitb_path",
                        nargs='*',
                        default=None)
    parser.add_argument("--vitl_path",
                        nargs='*',
                        default=None)
    parser.add_argument("--swint_path",
                        nargs='*',
                        default=None)
    parser.add_argument("--swinb_path",
                        nargs='*',
                        default=None)
    parser.add_argument("--fig_size",
                        nargs='*',
                        default=[10, 10])
    # parser.add_argument("--cam",
    #                     choices=['gradcam','scorecam'],
    #                     default='gradcam')

    args = parser.parse_args()

    return args
